# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: donggu

from experimental.prophet.contrib.functional import *


def test_vectorize():
  @vectorize
  def f(a, b, c, d):
    return a + b + c + d

  outs = f([1, 2], 0, 1, d=(3, 4, 4))
  assert outs == [[5, 6, 6], [6, 7, 7]]

  return 0
